#ifndef TEST_CNN_CIFAR_10
#define TEST_CNN_CIFAR_10

#include <torch/data/datasets/base.h>
#include <torch/data/example.h>
#include <torch/types.h>

#include <string>

/// The CIFAR-10 dataset.
class TORCH_API CIFAR10 : public torch::data::datasets::Dataset<CIFAR10> {
public:
  // cf.
  // https://discuss.pytorch.org/t/libtorch-c-equivalent-of-torch-transforms-compose-function-please-help/176605
  //
  // mean: 0.491400, 0.482158, 0.446531
  // std : 0.247032, 0.243485, 0.261588
  constexpr static std::array<double, 3> CIFAR10_MEANS = {0.491400, 0.482158,
                                                          0.446531};
  constexpr static std::array<double, 3> CIFAR10_STDS = {0.247032, 0.243485,
                                                         0.261588};

  /// The mode in which the dataset is loaded.
  enum class Mode { kTrain, kTest };

  /// Loads the CIFAR10 dataset from the `root` path.
  explicit CIFAR10(const std::string &root, Mode mode = Mode::kTrain,
                   bool check = false);

  /// Returns the `Example` at the given `index`.
  torch::data::Example<> get(size_t index) override;

  /// Returns the size of the dataset.
  torch::optional<size_t> size() const override;

  /// Returns true if this is the training subset of CIFAR10.
  // NOLINTNEXTLINE(bugprone-exception-escape)
  bool is_train() const noexcept;

  /// Returns all images stacked into a single tensor.
  const torch::Tensor &images() const;

  /// Returns all targets stacked into a single tensor.
  const torch::Tensor &targets() const;

private:
  torch::Tensor images_, targets_;
};

#endif /* TEST_CNN_CIFAR_10 */
